# -*- coding: utf-8 -*-

import numpy as np
import tensorflow as tf
from PIL import Image, ImageDraw, ImageFont
from tensorflow.python.framework import graph_io

from yolo_v3 import yolo_v3, load_weights, detections_boxes, non_max_suppression
import cv2

def load_coco_names(file_name):
    names = {}
    with open(file_name) as f:
        for id, name in enumerate(f):
            names[id] = name
    return names


def draw_boxes(font,boxes, img, cls_names, detection_size):
    names = {'person': '人', 'bicycle': '自行车', 'car': '汽车', 'motorbike': '摩托车', 'aeroplane': '飞机', 'bus': '公共汽车', 'train': '火车', 'truck': '卡车', 'boat': '船', 'traffic light': '交通灯', 'fire hydrant': '消火栓', 'stop sign': '停车标志', 'parking meter': '停车收费表', 'bench': '长凳', 'bird': '鸟', 'cat': '猫', 'dog': '狗', 'horse': '马', 'sheep': '羊', 'cow': '奶牛', 'elephant': '大象', 'bear': '熊', 'zebra': '斑马', 'giraffe': '长颈鹿', 'backpack': '背包', 'umbrella': '雨伞', 'handbag': '手提包', 'tie': '领带', 'suitcase': '手提箱', 'frisbee': '飞盘', 'skis': '滑雪板', 'snowboard': '滑雪板', 'sports ball': '运动球', 'kite': '风筝', 'baseball bat': '棒球棒', 'baseball glove': '棒球手套', 'skateboard': '滑板', 'surfboard': '冲浪板', 'tennis racket': '网球拍', 'bottle': '瓶子', 'wine glass': '酒杯', 'cup': '杯子', 'fork': '叉', 'knife': '刀', 'spoon': '勺子', 'bowl': '碗', 'banana': '香蕉', 'apple': '苹果', 'sandwich': '三明治', 'orange': '橙子', 'broccoli': '西兰花', 'carrot': '胡萝卜', 'hot dog': '热狗', 'pizza': '披萨', 'donut': '甜甜圈', 'cake': '蛋糕', 'chair': '椅子', 'sofa': '沙发', 'pottedplant': '盆栽植物', 'bed': '床', 'diningtable': '餐桌', 'toilet': '厕所', 'tvmonitor': '电视', 'laptop': '笔记本电脑', 'mouse': '鼠标', 'remote': '遥控器', 'keyboard': '键盘', 'cell phone': '手机', 'microwave': '微波炉', 'oven': '烤箱', 'toaster': '烤面包机', 'sink': '水池', 'refrigerator': '冰箱', 'book': '书', 'clock': '时钟', 'vase': '花瓶', 'scissors': '剪刀', 'teddy bear': '泰迪熊', 'hair drier': '吹风机', 'toothbrush': '牙刷'}
    draw = ImageDraw.Draw(img)

    for cls, bboxs in boxes.items():
        color = tuple(np.random.randint(0, 256, 3))
        for box, score in bboxs:
            box = convert_to_original_size(box, np.array(detection_size), np.array(img.size))
            draw.rectangle(box, outline=color)
            str1 = names[cls_names[cls].strip('\n')]
            draw.text(box[:2], str1, font=font, fill=color)
            #draw.text(box[:2], '{} {:.2f}%'.format(cls_names[cls], score * 100), fill=color)


def convert_to_original_size(box, size, original_size):
    ratio = original_size / size
    box = box.reshape(2, 2) * ratio
    return list(box.reshape(-1))


def tf_yolov3_init():
    global classes,inputs,detections,load_ops,boxes,sess
    classes = load_coco_names('../models/coco.names')  
    # placeholder for detector inputs
    inputs = tf.placeholder(tf.float32, [None, 416, 416, 3])

    with tf.variable_scope('detector'):
        detections = yolo_v3(inputs, len(classes), data_format='NHWC')
        load_ops = load_weights(tf.global_variables(scope='detector'), '../models/yolov3.weights')

    boxes = detections_boxes(detections)

    sess = tf.Session()
    sess.run(load_ops)  

def tf_yolov3_frame(font,frame,score):
    global classes,inputs,detections,load_ops,boxes,sess
    img_PIL = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
    img_resized = img_PIL.resize(size=(416, 416))
        
    detected_boxes = sess.run(boxes, feed_dict={inputs: [np.array(img_resized, dtype=np.float32)]})

    filtered_boxes = non_max_suppression(detected_boxes, confidence_threshold=score,
                                         iou_threshold=0.4)   
    draw_boxes(font,filtered_boxes, img_PIL, classes, (416, 416))
    frame = cv2.cvtColor(np.asarray(img_PIL),cv2.COLOR_RGB2BGR)
    return frame

def main():
    video_capture = cv2.VideoCapture(0) 
    tf_yolov3_init()
    font = ImageFont.truetype('NotoSansCJK-Thin.ttc', 30)
    while True:
        ret, frame = video_capture.read()
        frame=tf_yolov3_frame(font,frame)
        
        # Display the resulting image
        cv2.imshow('Video', frame)

        # Hit 'q' on the keyboard to quit!
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    # Release handle to the webcam
    sess.close()
    video_capture.release()
    cv2.destroyAllWindows()
   


if __name__ == '__main__':
    main()
